/*
 * This is written by Lulu Han and Yunguo Guan.
 * E-mail: locomotive_crypto@163.com
 */

#include <iostream>
#include <vector>
#include "she.h"
#include "timer.h"

void test_she_for_correctness(long k0, long k1, long k2);

void test_she_for_performance(long bits, long cnt);

int main(int argc, char* argv[]){
	std::cout << "----------------------------------------------------------" << std::endl;
	std::cout << "-----------------------<Usage>----------------------------" << std::endl;
	std::cout << "-----1) Run for correctness: ./program--------------------" << std::endl;
	std::cout << "-----2) Run for performance: ./program cnt----------------" << std::endl;
	std::cout << "----------------------------------------------------------" << std::endl << std::endl;
	
	
	if(argc == 1){
		test_she_for_correctness(1536, 32, 128);
	}
	else if(argc == 2){
		for(size_t i = 0; i < 10; ++i){
			test_she_for_performance(512, std::atoi(argv[1]));
		
			test_she_for_performance(1024, std::atoi(argv[1]));
		
			test_she_for_performance(1536, std::atoi(argv[1]));
			
			std::cout << std::endl << std::endl;
		}
	}
	
	return 0;
}


void test_she_for_correctness(long k0, long k1, long k2){
	long msg_len = 20;
	SHE::PrivateKey sk;
	SHE::PublicKey pk;
	
	SHE::key_gen(sk, pk, k0, k1, k2);
	
	SHE::Encryptor enc(pk);
	SHE::Decryptor dec(sk);
	
	NTL::ZZ pt, ct;
	pt = NTL::RandomLen_ZZ(msg_len);
	
	std::cout << "Testing for encryption and decryption..." << std::endl;
	
	std::cout << "pt = " << pt << std::endl;
	
	// Encrypt a message pt
	enc.encrypt(pt, ct);
	std::cout << "ct(pt) = " << ct << std::endl;
	
	// Decrypt a message ct
	NTL::ZZ rt;
	dec.decrypt(ct, rt);
	std::cout << "recover(pt) = " << rt << std::endl;
	
	// Testing for homomorphic addition
	std::cout << "\n1) Homomorphic addition for two ciphertexts..." << std::endl;
	NTL::ZZ pt2 , ct2;
	pt2 = NTL::RandomLen_ZZ(msg_len);
	
	std::cout << "pt2 = " << pt2 << std::endl;
	
	enc.encrypt(pt2, ct2);
	std::cout << "ct2(pt2) = " << ct2 << std::endl;
	
	enc.he_add_cc(ct, ct2, ct);
	std::cout << "ct(pt + pt2) = " << ct << std::endl;
	
	dec.decrypt(ct, rt);
	std::cout << "recover(pt + pt2) = " << rt << std::endl;
	
	if(rt != (pt + pt2)){
		throw std::runtime_error("Testing he_add_cp() functions failed...");
	}
	
	std::cout << "\n2) Homomorphic addition for one ciphertext and one plaintext..." << std::endl;
	NTL::ZZ pt3;
	pt3 = NTL::RandomLen_ZZ(msg_len);
	
	std::cout << "pt3 = " << pt3 << std::endl;
	
	enc.he_add_cp(ct2, pt3, ct);
	std::cout << "ct(pt2 + pt3) = " << ct << std::endl;
	
	dec.decrypt(ct, rt);
	std::cout << "recover(pt2 + pt3) = " << rt << std::endl;
	
	if(rt != (pt2 + pt3)){
		throw std::runtime_error("Testing he_add_cc() functions failed...");
	}
	
	std::cout << "\n3) Homomorphic multiplication for one ciphertext and one plaintext..." << std::endl;
	NTL::ZZ pt4;
	pt4 = NTL::RandomLen_ZZ(msg_len);
	
	std::cout << "pt4 = " << pt4 << std::endl;
	
	enc.he_mul_cp(ct2, pt4, ct);
	std::cout << "ct(pt2 * pt4) = " << ct << std::endl;
	
	dec.decrypt(ct, rt);
	std::cout << "recover(pt2 * pt4) = " << rt << std::endl;
	
	
	if(rt != (pt2 * pt4)){
		throw std::runtime_error("Testing he_mul_cp() functions failed...");
	}
	
	// Testing for homomorphic multiplication
	std::cout << "\n4) Homomorphic multiplication for one ciphertext and one ciphertext..." << std::endl;
	std::cout << "\tThe depth for homomorphic multiplication is: " << pk.get_mul_depth() << std::endl;
	NTL::ZZ ct4;
	
	enc.encrypt(pt4, ct4);
	std::cout << "ct(pt4) = " << ct4 << std::endl;
	
	enc.he_mul_cc(ct2, ct4, ct);
	std::cout << "ct(pt2 * pt4) = " << ct << std::endl;
	
	dec.decrypt(ct, rt);
	std::cout << "recover(pt2 * pt4) = " << rt << std::endl;
	
	if(rt != (pt2 * pt4)){
		throw std::runtime_error("Testing he_mul_cc() functions failed...");
	}
}

void test_she_for_performance(long bits, long cnt){
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	
	std::vector<SHE::Encryptor> venc;
	std::vector<SHE::Decryptor> vdec;
	
	long msg_len = 20;
	for(long i = 0; i < cnt; ++i){
		SHE::PrivateKey prikey;
		SHE::PublicKey pubkey;
	
		// Generate random key pair
		SHE::key_gen(prikey, pubkey, bits, 32, 128);
	
		// Define Encryptor and Decryptor objects
		SHE::Encryptor encryptor(pubkey);
		SHE::Decryptor decryptor(prikey);
		
		venc.push_back(encryptor);
		vdec.push_back(decryptor);
	}
	
	std::vector<NTL::ZZ> vpt, vct;
	NTL::ZZ pt, ct;
	for(long i = 0; i < cnt; ++i){
		pt = NTL::RandomLen_ZZ(msg_len);
	}
	
	time_point<high_resolution_clock> start;	
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		// Encrypt a message
		venc[i].encrypt(pt, ct);
		vct.push_back(ct);
	}
	double tm = time_from(start);
	std::cout << "The time for "<< cnt << " encryptions is : " << tm << std::endl;
	
	NTL::ZZ rt;
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		// Decrypt a ciphertext
		vdec[i].decrypt(vct[i], rt);
	}
	tm = time_from(start);
	std::cout << "The time for "<< cnt << " decryptions is : " << tm << std::endl;
	
	std::vector<NTL::ZZ> vpt2, vct2;
	NTL::ZZ pt2, ct2;
	for(long i = 0; i < cnt; ++i){
		pt2 = NTL::RandomLen_ZZ(msg_len);
		vpt2.push_back(pt2);
		
		venc[i].encrypt(pt2, ct2);
		vct2.push_back(ct2);
	}
	
	// 1) Homomorphic addition for two ciphertexts
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		venc[i].he_add_cc(vct[i], vct2[i], ct);
	}
	tm = time_from(start);
	std::cout << "The time for "<< cnt << " homomorphic addition for two ciphertexts is : " << tm << std::endl;
	
	// 2) Homomorphic addition for one ciphertext and one plaintext
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		venc[i].he_add_cp(vct[i], vpt2[i], ct);
	}
	tm = time_from(start);
	std::cout << "The time for "<< cnt << " homomorphic addition for one ciphertext and one plaintext is : " << tm << std::endl;
	
	// 3) Homomorphic multiplication for two ciphertexts
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		venc[i].he_mul_cc(vct[i], vct2[i], ct);
	}
	tm = time_from(start);
	std::cout << "The time for "<< cnt << " homomorphic multiplication for two ciphertexts is : " << tm << std::endl;
	
	// 4) Homomorphic multiplication for one ciphertext and one plaintext
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		venc[i].he_mul_cp(vct[i], vpt2[i], ct);
	}
	tm = time_from(start);
	std::cout << "The time for "<< cnt << " homomorphic multiplication for one ciphertext and one plaintext is : " << tm << std::endl;
}
